"""
多模型融合运行脚本
Author: Model Team
Date: 2025-11-16
Description: Ensemble model combining multiple machine learning algorithms
"""

import os
import sys
import json
import pandas as pd
import numpy as np
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.linear_model import Ridge, LinearRegression
from sklearn.ensemble import VotingRegressor
import joblib
import warnings
from datetime import datetime
from scipy import stats
import matplotlib.pyplot as plt
import seaborn as sns

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from model.validation_report_utils import generate_validation_outputs

warnings.filterwarnings('ignore')

class MultiModelEnsemble:
    def __init__(self, config_path='model_config.json'):
        """初始化多模型融合器"""
        # 确保我们在正确的目录中运行
        script_dir = os.path.dirname(os.path.abspath(__file__))
        os.chdir(script_dir)

        self.config = self.load_config(config_path)
        self.base_models = {}
        self.ensemble_models = {}
        self.X_train = None
        self.y_train = None
        self.X_test = None
        self.test_ids = None
        self.feature_names = None
        self.ensemble_results = {}
        self.best_method = None

        # 设置输出目录
        self.setup_output_directories()

    def setup_output_directories(self):
        """设置输出目录结构"""
        self.output_dirs = {
            'models': 'outputs/models',
            'submissions': 'outputs/submissions',
            'reports': 'outputs/reports',
            'logs': 'outputs/logs'
        }

        # 创建所有输出目录
        for dir_name, dir_path in self.output_dirs.items():
            os.makedirs(dir_path, exist_ok=True)
            print(f"[OK] 创建目录: {dir_path}")

    def load_config(self, config_path):
        """加载模型配置"""
        try:
            if not os.path.exists(config_path):
                print(f"[ERROR] 配置文件不存在: {config_path}")
                sys.exit(1)

            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)

            print(f"[OK] 配置文件加载成功")
            print(f"模型名称: {config['ensemble_info']['ensemble_name']}")
            return config

        except Exception as e:
            print(f"[ERROR] 配置文件加载失败: {e}")
            sys.exit(1)

    def load_data(self):
        """加载数据"""
        print("[INFO] 加载数据...")

        try:
            # 获取项目根目录
            current_dir = os.getcwd()
            script_dir = os.path.dirname(os.path.abspath(__file__))
            model_dir = os.path.dirname(script_dir)
            project_root = os.path.dirname(model_dir)

            # 加载数据
            train_path = os.path.join(project_root, 'feature/train_feature.csv')
            test_path = os.path.join(project_root, 'feature/test_feature.csv')

            if not os.path.exists(train_path) or not os.path.exists(test_path):
                print(f"[ERROR] 数据文件不存在: {train_path} 或 {test_path}")
                return False

            train_data = pd.read_csv(train_path)
            test_data = pd.read_csv(test_path)

            # 准备训练数据
            self.X_train = train_data.drop(['SalePrice', 'Id'], axis=1, errors='ignore')
            self.y_train = np.log1p(train_data['SalePrice'])  # 对数变换
            self.X_test = test_data.drop(['Id'], axis=1, errors='ignore')
            self.test_ids = test_data['Id'] if 'Id' in test_data.columns else None

            self.feature_names = list(self.X_train.columns)

            print(f"[OK] 数据加载成功")
            print(f"   训练集形状: {self.X_train.shape}")
            print(f"   测试集形状: {self.X_test.shape}")
            print(f"   特征数量: {len(self.feature_names)}")
            return True

        except Exception as e:
            print(f"[ERROR] 数据加载失败: {e}")
            return False

    def load_base_models(self):
        """加载基础模型"""
        print("[INFO] 加载基础模型...")

        base_models_config = self.config['base_models']
        loaded_models = 0

        for model_name, model_config in base_models_config.items():
            if not model_config['enabled']:
                continue

            try:
                print(f"[INFO] 加载 {model_name} 模型...")

                # 加载模型
                model_path = os.path.join(os.getcwd(), model_config['model_path'])
                if os.path.exists(model_path):
                    model = joblib.load(model_path)
                    self.base_models[model_name] = {
                        'model': model,
                        'weight': model_config['weight'],
                        'config': model_config
                    }
                    loaded_models += 1
                    print(f"[OK] {model_name} 模型加载成功")
                else:
                    print(f"[WARNING] {model_name} 模型文件不存在: {model_path}")

                # 加载特征名称（如果存在）
                if 'feature_names_path' in model_config:
                    feature_path = os.path.join(os.getcwd(), model_config['feature_names_path'])
                    if os.path.exists(feature_path):
                        feature_names = joblib.load(feature_path)
                        # 确保特征名称一致
                        if set(feature_names) == set(self.feature_names):
                            print(f"[OK] {model_name} 特征名称匹配")
                        else:
                            print(f"[WARNING] {model_name} 特征名称不匹配")
                    else:
                        print(f"[WARNING] {model_name} 特征名称文件不存在")

            except Exception as e:
                print(f"[ERROR] 加载 {model_name} 模型失败: {e}")

        print(f"[OK] 成功加载 {loaded_models} 个基础模型")
        return loaded_models > 0

    def create_simple_average_ensemble(self):
        """创建简单平均融合模型"""
        print("[INFO] 创建简单平均融合模型...")

        def ensemble_predict(X):
            predictions = []
            for model_name, model_info in self.base_models.items():
                pred = model_info['model'].predict(X)
                predictions.append(pred)
            return np.mean(predictions, axis=0)

        self.ensemble_models['simple_average'] = {
            'predict': ensemble_predict,
            'description': 'Simple average of all model predictions'
        }
        print("[OK] 简单平均融合模型创建成功")

    def create_weighted_average_ensemble(self):
        """创建加权平均融合模型"""
        print("[INFO] 创建加权平均融合模型...")

        def ensemble_predict(X):
            predictions = []
            weights = []
            for model_name, model_info in self.base_models.items():
                pred = model_info['model'].predict(X)
                predictions.append(pred)
                weights.append(model_info['weight'])

            # 归一化权重
            weights = np.array(weights) / np.sum(weights)
            weighted_pred = np.average(predictions, axis=0, weights=weights)
            return weighted_pred

        self.ensemble_models['weighted_average'] = {
            'predict': ensemble_predict,
            'description': 'Weighted average using model weights'
        }
        print("[OK] 加权平均融合模型创建成功")

    def create_voting_ensemble(self):
        """创建投票融合模型"""
        print("[INFO] 创建投票融合模型...")

        # 准备模型列表用于VotingRegressor
        estimators = []
        for model_name, model_info in self.base_models.items():
            estimators.append((model_name, model_info['model']))

        # 创建投票融合模型
        voting_model = VotingRegressor(estimators=estimators)
        voting_model.fit(self.X_train, self.y_train)

        self.ensemble_models['voting_regressor'] = {
            'model': voting_model,
            'predict': voting_model.predict,
            'description': 'Sklearn VotingRegressor ensemble'
        }
        print("[OK] 投票融合模型创建成功")

    def create_stacking_ensemble(self):
        """创建堆叠融合模型"""
        print("[INFO] 创建堆叠融合模型...")

        stacking_config = self.config['ensemble_methods']['stacking']
        cv_folds = stacking_config['cv_folds']
        meta_learner_type = stacking_config['meta_learner']

        # 准备基础模型的预测结果作为元特征
        kf = KFold(n_splits=cv_folds, shuffle=True, random_state=42)

        # 为训练集生成基础模型的预测
        meta_features_train = np.zeros((len(self.X_train), len(self.base_models)))
        meta_features_test = np.zeros((len(self.X_test), len(self.base_models)))

        for i, (model_name, model_info) in enumerate(self.base_models.items()):
            print(f"[INFO] 生成 {model_name} 的堆叠特征...")

            # 交叉验证预测
            meta_features_train[:, i] = self._get_cv_predictions(
                model_info['model'], self.X_train, self.y_train, kf
            )

            # 测试集预测
            model_info['model'].fit(self.X_train, self.y_train)
            meta_features_test[:, i] = model_info['model'].predict(self.X_test)

        # 选择元学习器
        if meta_learner_type == 'ridge':
            meta_learner = Ridge(alpha=1.0, random_state=42)
        elif meta_learner_type == 'linear':
            meta_learner = LinearRegression()
        else:
            meta_learner = Ridge(alpha=1.0, random_state=42)

        # 训练元学习器
        meta_learner.fit(meta_features_train, self.y_train)

        def stacking_predict(X):
            # 获取基础模型预测
            meta_features = np.zeros((len(X), len(self.base_models)))
            for i, model_info in enumerate(self.base_models.values()):
                meta_features[:, i] = model_info['model'].predict(X)
            return meta_learner.predict(meta_features)

        self.ensemble_models['stacking'] = {
            'meta_learner': meta_learner,
            'predict': stacking_predict,
            'description': f'Stacking with {meta_learner_type} meta-learner'
        }
        print("[OK] 堆叠融合模型创建成功")

    def _get_cv_predictions(self, model, X, y, kf):
        """获取交叉验证预测结果"""
        predictions = np.zeros(len(X))

        for train_idx, val_idx in kf.split(X):
            X_train_fold, X_val_fold = X.iloc[train_idx], X.iloc[val_idx]
            y_train_fold = y.iloc[train_idx]

            # 训练模型
            model.fit(X_train_fold, y_train_fold)

            # 预测
            pred = model.predict(X_val_fold)
            predictions[val_idx] = pred

        return predictions

    def evaluate_ensemble_methods(self):
        """评估所有融合方法"""
        print("[INFO] 评估融合方法...")

        if not self.config['training_params']['compare_methods']:
            print("[INFO] 跳过方法比较")
            return

        cv_config = self.config['data_params']
        kf = KFold(n_splits=cv_config['cv_splits'], shuffle=cv_config['shuffle'], random_state=42)

        method_results = {}

        for method_name, ensemble_info in self.ensemble_models.items():
            print(f"[INFO] 评估 {method_name}...")

            rmse_scores = []
            mae_scores = []
            r2_scores = []

            for fold, (train_idx, val_idx) in enumerate(kf.split(self.X_train), 1):
                X_train_fold, X_val_fold = self.X_train.iloc[train_idx], self.X_train.iloc[val_idx]
                y_train_fold, y_val_fold = self.y_train.iloc[train_idx], self.y_train.iloc[val_idx]

                # 为融合模型准备数据
                if method_name == 'stacking':
                    # 堆叠模型需要重新训练基础模型
                    self._train_base_models_for_fold(X_train_fold, y_train_fold)

                # 预测
                y_pred = ensemble_info['predict'](X_val_fold)

                # 计算指标
                rmse = np.sqrt(mean_squared_error(y_val_fold, y_pred))
                mae = mean_absolute_error(y_val_fold, y_pred)
                r2 = r2_score(y_val_fold, y_pred)

                rmse_scores.append(rmse)
                mae_scores.append(mae)
                r2_scores.append(r2)

                print(f"第 {fold} 折完成 - RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}")

            # 记录结果
            method_results[method_name] = {
                'RMSE': {
                    'mean': float(np.mean(rmse_scores)),
                    'std': float(np.std(rmse_scores)),
                    'folds': [float(score) for score in rmse_scores]
                },
                'MAE': {
                    'mean': float(np.mean(mae_scores)),
                    'std': float(np.std(mae_scores)),
                    'folds': [float(score) for score in mae_scores]
                },
                'R2': {
                    'mean': float(np.mean(r2_scores)),
                    'std': float(np.std(r2_scores)),
                    'folds': [float(score) for score in r2_scores]
                }
            }

            print(f"{method_name} 平均 RMSE: {method_results[method_name]['RMSE']['mean']:.4f} ± {method_results[method_name]['RMSE']['std']:.4f}")

        self.ensemble_results = method_results

        # 选择最佳方法
        if self.config['training_params']['select_best_method']:
            best_method = min(method_results.keys(),
                            key=lambda x: method_results[x]['RMSE']['mean'])
            self.best_method = best_method
            print(f"[INFO] 最佳融合方法: {best_method} (RMSE: {method_results[best_method]['RMSE']['mean']:.4f})")

    def _train_base_models_for_fold(self, X_train, y_train):
        """为堆叠融合训练基础模型"""
        for model_info in self.base_models.values():
            model_info['model'].fit(X_train, y_train)

    def generate_predictions(self):
        """生成预测"""
        print("[INFO] 生成预测...")

        try:
            # 确保所有基础模型都在完整数据上训练过
            for model_name, model_info in self.base_models.items():
                model_info['model'].fit(self.X_train, self.y_train)

            for method_name, ensemble_info in self.ensemble_models.items():
                print(f"[INFO] 生成 {method_name} 预测...")

                # 预测并还原对数变换
                y_pred_log = ensemble_info['predict'](self.X_test)
                y_pred = np.expm1(y_pred_log)

                # 生成提交文件
                submission = pd.DataFrame({
                    'Id': self.test_ids,
                    'SalePrice': y_pred
                })

                submission_filename = f"submission_{method_name}.csv"
                submission_path = os.path.join(os.getcwd(), self.output_dirs['submissions'], submission_filename)
                os.makedirs(os.path.dirname(submission_path), exist_ok=True)
                submission.to_csv(submission_path, index=False)

                print(f"[OK] {method_name} 提交文件生成成功: {submission_path}")
                print(f"   预测结果统计:")
                print(f"   最小值: ${y_pred.min():,.2f}")
                print(f"   最大值: ${y_pred.max():,.2f}")
                print(f"   平均值: ${y_pred.mean():,.2f}")

            return True

        except Exception as e:
            print(f"[ERROR] 预测生成失败: {e}")
            return False

    def save_models(self):
        """保存融合模型"""
        print("[INFO] 保存融合模型...")

        try:
            for method_name, ensemble_info in self.ensemble_models.items():
                if 'model' in ensemble_info:
                    # 保存sklearn模型
                    model_filename = f"{method_name}_ensemble_model.pkl"
                    model_path = os.path.join(os.getcwd(), self.output_dirs['models'], model_filename)
                    joblib.dump(ensemble_info['model'], model_path)
                    print(f"[OK] {method_name} 模型保存成功: {model_path}")

                if 'meta_learner' in ensemble_info:
                    # 保存元学习器
                    meta_filename = f"{method_name}_meta_learner.pkl"
                    meta_path = os.path.join(os.getcwd(), self.output_dirs['models'], meta_filename)
                    joblib.dump(ensemble_info['meta_learner'], meta_path)
                    print(f"[OK] {method_name} 元学习器保存成功: {meta_path}")

            # 保存融合结果
            results_filename = "ensemble_results.json"
            results_path = os.path.join(os.getcwd(), self.output_dirs['models'], results_filename)

            save_results = {
                'ensemble_info': self.config['ensemble_info'],
                'base_models': list(self.base_models.keys()),
                'ensemble_methods': self.ensemble_results,
                'best_method': self.best_method
            }

            with open(results_path, 'w', encoding='utf-8') as f:
                json.dump(save_results, f, indent=2, ensure_ascii=False)

            print(f"[OK] 融合结果保存成功: {results_path}")
            return True

        except Exception as e:
            print(f"[ERROR] 模型保存失败: {e}")
            return False

    def generate_reports(self):
        """生成报告"""
        print("[INFO] 生成分析报告...")

        try:
            if not self.ensemble_results:
                print("[WARNING] 没有评估结果，跳过报告生成")
                return False

            # 生成方法对比图
            if self.config['output_params']['plot_comparison']:
                self._plot_method_comparison()

            # 生成最佳方法的验证报告
            if self.best_method and self.config['output_params']['generate_report']:
                self._generate_best_method_report()

            print("[OK] 报告生成完成")
            return True

        except Exception as e:
            print(f"[ERROR] 报告生成失败: {e}")
            return False

    def _plot_method_comparison(self):
        """绘制方法对比图"""
        print("[INFO] 生成方法对比图...")

        methods = list(self.ensemble_results.keys())
        rmse_means = [self.ensemble_results[method]['RMSE']['mean'] for method in methods]
        rmse_stds = [self.ensemble_results[method]['RMSE']['std'] for method in methods]

        plt.figure(figsize=(12, 8))

        # RMSE对比
        plt.subplot(2, 2, 1)
        bars = plt.bar(methods, rmse_means, yerr=rmse_stds, capsize=5, alpha=0.7)
        plt.ylabel('RMSE')
        plt.title('Ensemble Methods RMSE Comparison')
        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3)

        # 标记最佳方法
        best_idx = np.argmin(rmse_means)
        bars[best_idx].set_color('red')
        plt.legend(['Methods', f'Best: {methods[best_idx]}'])

        # MAE对比
        mae_means = [self.ensemble_results[method]['MAE']['mean'] for method in methods]
        mae_stds = [self.ensemble_results[method]['MAE']['std'] for method in methods]

        plt.subplot(2, 2, 2)
        bars = plt.bar(methods, mae_means, yerr=mae_stds, capsize=5, alpha=0.7, color='orange')
        plt.ylabel('MAE')
        plt.title('Ensemble Methods MAE Comparison')
        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3)
        bars[best_idx].set_color('red')

        # R²对比
        r2_means = [self.ensemble_results[method]['R2']['mean'] for method in methods]
        r2_stds = [self.ensemble_results[method]['R2']['std'] for method in methods]

        plt.subplot(2, 2, 3)
        bars = plt.bar(methods, r2_means, yerr=r2_stds, capsize=5, alpha=0.7, color='green')
        plt.ylabel('R²')
        plt.title('Ensemble Methods R² Comparison')
        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3)
        bars[best_idx].set_color('red')

        # 综合对比
        plt.subplot(2, 2, 4)
        # 归一化指标进行比较
        normalized_rmse = np.array(rmse_means) / np.max(rmse_means)
        normalized_mae = np.array(mae_means) / np.max(mae_means)
        normalized_r2 = np.array(r2_means) / np.max(r2_means)

        x = np.arange(len(methods))
        width = 0.25

        plt.bar(x - width, normalized_rmse, width, label='RMSE', alpha=0.7)
        plt.bar(x, normalized_mae, width, label='MAE', alpha=0.7)
        plt.bar(x + width, normalized_r2, width, label='R²', alpha=0.7)

        plt.xlabel('Ensemble Methods')
        plt.ylabel('Normalized Score')
        plt.title('Normalized Performance Comparison')
        plt.xticks(x, methods, rotation=45)
        plt.legend()
        plt.grid(True, alpha=0.3)

        plt.tight_layout()

        # 保存图片
        plot_path = os.path.join(os.getcwd(), self.output_dirs['reports'], 'ensemble_comparison.png')
        plt.savefig(plot_path, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"[OK] 方法对比图保存成功: {plot_path}")

    def _generate_best_method_report(self):
        """生成最佳方法的验证报告"""
        print(f"[INFO] 生成最佳方法 {self.best_method} 的验证报告...")

        # 重新获取最佳方法的详细fold结果
        cv_config = self.config['data_params']
        kf = KFold(n_splits=cv_config['cv_splits'], shuffle=cv_config['shuffle'], random_state=42)

        fold_results = []
        ensemble_info = self.ensemble_models[self.best_method]

        for fold, (train_idx, val_idx) in enumerate(kf.split(self.X_train), 1):
            X_train_fold, X_val_fold = self.X_train.iloc[train_idx], self.X_train.iloc[val_idx]
            y_train_fold, y_val_fold = self.y_train.iloc[train_idx], self.y_train.iloc[val_idx]

            # 重新训练（如果需要）
            if self.best_method == 'stacking':
                self._train_base_models_for_fold(X_train_fold, y_train_fold)

            # 预测
            y_pred = ensemble_info['predict'](X_val_fold)

            # 收集fold结果
            for i, (true_val, pred_val) in enumerate(zip(y_val_fold, y_pred)):
                fold_results.append({
                    'fold': fold,
                    'Id': int(val_idx[i]),
                    'y_true': float(true_val),
                    'y_pred': float(pred_val)
                })

        # 生成报告
        try:
            # 临时切换到项目根目录
            original_cwd = os.getcwd()
            script_dir = os.path.dirname(os.path.abspath(__file__))
            model_dir = os.path.dirname(script_dir)
            project_root = os.path.dirname(model_dir)
            os.chdir(project_root)

            generate_validation_outputs(
                folds_results=fold_results,
                metrics_dict=self.ensemble_results[self.best_method],
                model_name=self.config['ensemble_info']['ensemble_name'],
                model_version=self.config['ensemble_info']['version'],
                framework=f"{self.best_method}_ensemble",
                target=self.config['data_params']['target_variable'],
                train_size=len(self.X_train),
                n_features=len(self.feature_names),
                n_splits=cv_config['cv_splits'],
                random_state=42,
                training_time="Ensemble training time",
                notes=f"Best ensemble method: {self.best_method}. {ensemble_info['description']}",
                args={
                    'base_models': list(self.base_models.keys()),
                    'ensemble_method': self.best_method,
                    'ensemble_info': self.config['ensemble_info']
                }
            )

            # 恢复工作目录
            os.chdir(original_cwd)

            # 移动验证报告到outputs/reports目录
            validation_files = [
                f'validation_{self.best_method}_ensemble.csv',
                f'validation_report_{self.best_method}_ensemble.json'
            ]

            for file_name in validation_files:
                src_path = os.path.join(project_root, file_name)
                dst_path = os.path.join(self.output_dirs['reports'], file_name)
                if os.path.exists(src_path):
                    import shutil
                    shutil.move(src_path, dst_path)
                    print(f"[OK] 移动验证报告到: {dst_path}")

        except Exception as e:
            print(f"[WARNING] 最佳方法验证报告生成失败: {e}")

    def run_training(self):
        """运行完整的训练流程"""
        print(f"[INFO] 开始训练模型: {self.config['ensemble_info']['ensemble_name']}")
        print("="*60)

        steps = [
            ("数据加载", self.load_data),
            ("加载基础模型", self.load_base_models),
        ]

        # 执行基础步骤
        for step_name, step_func in steps:
            print(f"\n[INFO] {step_name}...")
            if not step_func():
                print(f"[ERROR] {step_name}失败")
                return False

        # 创建融合模型
        print("\n[INFO] 创建融合模型...")
        if self.config['ensemble_methods']['simple_average']['enabled']:
            self.create_simple_average_ensemble()

        if self.config['ensemble_methods']['weighted_average']['enabled']:
            self.create_weighted_average_ensemble()

        if self.config['ensemble_methods']['voting_regressor']['enabled']:
            self.create_voting_ensemble()

        if self.config['ensemble_methods']['stacking']['enabled']:
            self.create_stacking_ensemble()

        # 评估融合方法
        self.evaluate_ensemble_methods()

        # 执行后续步骤
        final_steps = [
            ("保存模型", self.save_models),
            ("生成预测", self.generate_predictions),
            ("生成报告", self.generate_reports)
        ]

        for step_name, step_func in final_steps:
            print(f"\n[INFO] {step_name}...")
            if not step_func():
                print(f"[ERROR] {step_name}失败")
                return False

        return True

def main():
    """主函数"""
    try:
        runner = MultiModelEnsemble()
        success = runner.run_training()

        if success:
            print("\n" + "="*60)
            print("🎉 多模型融合训练完成!")
            print("="*60)
            print("📁 输出文件:")
            for method_name in runner.ensemble_models.keys():
                print(f"   - {method_name} 提交文件: submission_{method_name}.csv")
            print(f"   - 融合结果: ensemble_results.json")
            print(f"   - 方法对比图: ensemble_comparison.png")
            if runner.best_method:
                print(f"   - 最佳方法: {runner.best_method}")
            print(f"   - 所有文件已整理到 outputs/ 目录")
        else:
            print("\n[ERROR] 模型融合训练失败!")
            sys.exit(1)

    except Exception as e:
        print(f"\n[ERROR] 运行过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()